import sqlite3
import json
import csv
import os
from datetime import datetime, timedelta


class DataAnalyzer:
    """数据分析器"""

    def __init__(self, db_path="work_monitor.db"):
        self.db_path = db_path
        self.init_database()

    def init_database(self):
        """初始化数据库"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            # 创建工作数据表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS work_data (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    keystrokes INTEGER,
                    mouse_distance REAL,
                    total_work_time REAL,
                    applications TEXT
                )
            ''')

            # 创建网络数据表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS network_data (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    connections TEXT
                )
            ''')

            conn.commit()
            conn.close()
            print("数据分析数据库初始化成功")
        except Exception as e:
            print(f"数据分析数据库初始化失败: {e}")

    def save_work_data(self, work_data):
        """保存工作数据"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            cursor.execute('''
                INSERT INTO work_data (keystrokes, mouse_distance, total_work_time, applications)
                VALUES (?, ?, ?, ?)
            ''', (
                work_data['keystrokes'],
                work_data['mouse_distance'],
                work_data['total_work_time'],
                json.dumps(work_data['applications'])
            ))

            conn.commit()
            conn.close()
        except Exception as e:
            print(f"保存工作数据失败: {e}")

    def save_network_data(self, network_connections):
        """保存网络数据"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            cursor.execute('''
                INSERT INTO network_data (connections)
                VALUES (?)
            ''', (json.dumps(network_connections),))

            conn.commit()
            conn.close()
        except Exception as e:
            print(f"保存网络数据失败: {e}")

    def calculate_productivity_score(self, work_data):
        """计算生产力指数"""
        try:
            work_time_hours = work_data['total_work_time'] / 3600
            if work_time_hours == 0:
                return 0

            # 基于按键和鼠标活动计算生产力
            keyboard_activity = min(work_data['keystrokes'] / 1000, 1)
            mouse_activity = min(work_data['mouse_distance'] * 0.0264583 / 1000, 1)

            # 应用数量因子
            app_factor = min(len(work_data['applications']) / 5, 1)

            # 综合计算
            productivity = int((keyboard_activity + mouse_activity + app_factor) * 33.33)
            return min(productivity, 100)
        except Exception as e:
            print(f"计算生产力指数失败: {e}")
            return 0

    def export_data(self, format_type='csv', days=30, save_dir=None):
        """导出数据"""
        try:
            if save_dir is None:
                save_dir = os.path.expanduser("~/Desktop")

            # 计算日期范围
            end_date = datetime.now()
            start_date = end_date - timedelta(days=days)

            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()

            result_files = []

            if format_type == 'csv':
                # 导出工作数据
                work_file = os.path.join(save_dir, f"work_data_{end_date.strftime('%Y%m%d')}.csv")
                cursor.execute('''
                    SELECT * FROM work_data 
                    WHERE timestamp >= ? AND timestamp <= ?
                    ORDER BY timestamp DESC
                ''', (start_date, end_date))

                work_data = cursor.fetchall()

                with open(work_file, 'w', newline='', encoding='utf-8') as f:
                    writer = csv.writer(f)
                    writer.writerow(['ID', '时间戳', '按键次数', '鼠标距离', '工作时长', '应用程序'])
                    writer.writerows(work_data)

                result_files.append(work_file)

                # 导出网络数据
                network_file = os.path.join(save_dir, f"network_data_{end_date.strftime('%Y%m%d')}.csv")
                cursor.execute('''
                    SELECT * FROM network_data 
                    WHERE timestamp >= ? AND timestamp <= ?
                    ORDER BY timestamp DESC
                ''', (start_date, end_date))

                network_data = cursor.fetchall()

                with open(network_file, 'w', newline='', encoding='utf-8') as f:
                    writer = csv.writer(f)
                    writer.writerow(['ID', '时间戳', '网络连接'])
                    writer.writerows(network_data)

                result_files.append(network_file)

            conn.close()
            return result_files

        except Exception as e:
            print(f"导出数据失败: {e}")
            return []
